package com.gjy.encryption.asymmetric;

import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.spec.SecretKeySpec;
import java.security.*;
import java.security.interfaces.ECPrivateKey;
import java.security.interfaces.ECPublicKey;
import java.security.spec.ECParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

/**
 * DH 加解密组件
 * 需要注意的是, 甲方的本地秘钥是由甲方的私钥和乙方的公钥生成,
 * 乙方的本地秘钥是由乙方的私钥和甲方的公钥生成, 可参考示例代码
 * <p>
 * 示例代码:
 * ECDHCoder coder = new ECDHCoder();
 * Map<String, Key> m1 = coder.initKey();
 * byte[] pu1 = m1.get(ECDHCoder.PUBLIC_KEY).getEncoded();
 * byte[] pr1 = m1.get(ECDHCoder.PRIVATE_KEY).getEncoded();
 * <p>
 * Map<String, Key> m2 = coder.initKey(pu1);
 * byte[] pu2 = m2.get(ECDHCoder.PUBLIC_KEY).getEncoded();
 * byte[] pr2 = m2.get(ECDHCoder.PRIVATE_KEY).getEncoded();
 * <p>
 * byte[] a = coder.getSecretKey(pu2, pr1);
 * byte[] b = coder.getSecretKey(pu1, pr2);
 * <p>
 * byte[] encrypt = coder.encrypt("Hello".getBytes(), a);
 * <p>
 * byte[] decrypt = coder.decrypt(encrypt, b);
 * <p>
 * System.out.println(new String(decrypt));
 *
 * @author 宫静雨
 * @version 1.0
 * @since 2022-11-07 12:35:40
 */
public class ECDHCoder {

    public ECDHCoder() {
        Security.addProvider(new BouncyCastleProvider());
    }

    // 非对称加解密算法
    private static final String KEY_ALGORITHM = "ECDH";
    // 本地秘钥算法，及对称加密算法，可选：Blowfish, RC2, RC4
    private static final String SECRET_ALGORITHM = "Blowfish";
    // 秘钥长度，范围112-571之间，默认256
    private static final Integer KEY_SIZE = 256;
    public static final String PUBLIC_KEY = "ECDHPublicKey";
    public static final String PRIVATE_KEY = "ECDHPrivateKey";

    /**
     * 加密
     *
     * @param data 待加密数据
     * @param key  秘钥
     * @return 加密数据
     */
    public byte[] encrypt(byte[] data, byte[] key) {
        try {
            // 生成本地秘钥
            SecretKeySpec secretKeySpec = new SecretKeySpec(key, SECRET_ALGORITHM);
            Cipher cipher = Cipher.getInstance(secretKeySpec.getAlgorithm());
            cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec);
            return cipher.doFinal(data);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 解密
     *
     * @param data 待解密数据
     * @param key  秘钥
     * @return 解密数据
     */
    public byte[] decrypt(byte[] data, byte[] key) {
        try {
            // 生成本地秘钥
            SecretKeySpec secretKeySpec = new SecretKeySpec(key, SECRET_ALGORITHM);
            Cipher cipher = Cipher.getInstance(secretKeySpec.getAlgorithm());
            cipher.init(Cipher.DECRYPT_MODE, secretKeySpec);
            return cipher.doFinal(data);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 构建本地秘钥
     *
     * @param publicKey  公钥
     * @param privateKey 私钥
     * @return 本地秘钥
     */
    public byte[] getSecretKey(byte[] publicKey, byte[] privateKey) {
        try {
            // 实例化秘钥工厂
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            // 初始化公钥，秘钥材料转换
            X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(publicKey);
            // 产生公钥
            PublicKey pubKey = keyFactory.generatePublic(x509EncodedKeySpec);

            // 初始化私钥，秘钥转换材料
            PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(privateKey);
            // 产生私钥
            PrivateKey priKey = keyFactory.generatePrivate(pkcs8EncodedKeySpec);

            // 实例化
            KeyAgreement keyAgreement = KeyAgreement.getInstance(keyFactory.getAlgorithm());
            keyAgreement.init(priKey);
            keyAgreement.doPhase(pubKey, true);
            return keyAgreement.generateSecret(SECRET_ALGORITHM).getEncoded();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 初始化甲方秘钥
     */
    public Map<String, Key> initKey() {
        try {
            // 实例化密钥对生成器
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM);
            keyPairGenerator.initialize(KEY_SIZE);
            // 生成秘钥对
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            ECPublicKey publicKey = (ECPublicKey) keyPair.getPublic();
            ECPrivateKey privateKey = (ECPrivateKey) keyPair.getPrivate();

            Map<String, Key> PARTY_A_MAP = new HashMap<>(4);
            PARTY_A_MAP.put(PUBLIC_KEY, publicKey);
            PARTY_A_MAP.put(PRIVATE_KEY, privateKey);
            return PARTY_A_MAP;
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 初始化乙方秘钥
     *
     * @param partyAPublicKey 甲方公钥
     */
    public Map<String, Key> initKey(byte[] partyAPublicKey) {
        try {
            // 解析甲方公钥，转换公钥材料
            X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(partyAPublicKey);
            // 实例化工厂
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            // 产生公钥
            PublicKey publicKey = keyFactory.generatePublic(x509EncodedKeySpec);
            // 由甲方公钥构建乙方秘钥
            ECParameterSpec dhParameterSpec = ((ECPublicKey) publicKey).getParams();

            // 实例化密钥对生成器
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(keyFactory.getAlgorithm());
            // 初始化密钥对生成器
            keyPairGenerator.initialize(dhParameterSpec);
            // 产生密钥对
            KeyPair keyPair = keyPairGenerator.genKeyPair();
            ECPublicKey dhPublicKey = (ECPublicKey) keyPair.getPublic();
            ECPrivateKey dhPrivateKey = (ECPrivateKey) keyPair.getPrivate();

            Map<String, Key> PARTY_B_MAP = new HashMap<>(4);
            PARTY_B_MAP.put(PUBLIC_KEY, dhPublicKey);
            PARTY_B_MAP.put(PRIVATE_KEY, dhPrivateKey);
            return PARTY_B_MAP;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
